# -*- coding:utf-8 -*-
import traceback

# @Time    : 2023/7/9 23:16
# @Author  : zengwenjia
# @Email   : zengwenjia@lingxi.ai
# @File    : draw_picture.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import networkx as nx

import settings

# 为matplotlib设置中文字体路径
plt.rcParams['font.sans-serif'] = ['SimHei']  # SimHei是黑体
if not settings.ONLINE:
    plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题


class DrawPicture():
    def __init__(self):
        pass

    @classmethod
    def save_financial_trend(cls, data, path):
        try:
            if not data.get('detailsVos', None):
                return False
            else:
                # 将dict的list转成pandas
                df = pd.DataFrame(data['detailsVos'])
                x = df["age"]
                if 'pensionPrepare' in str(data):
                    y1 = df["pensionPrepare"]
                if 'livingCost' in str(data):
                    y2 = df["livingCost"]

                fig, ax = plt.subplots()
                if 'pensionPrepare' in str(data):
                    ax.plot(x, y1, color='#e18d4d', label='净收入',
                            linewidth=1.2)
                if 'livingCost' in str(data):
                    ax.plot(x, y2, color='#6fa7f8', label='净支出',
                            linewidth=1.2)

                # 使用指定的颜色在相同的位置绘制散点
                if 'pensionPrepare' in str(data):
                    ax.scatter(x, y1, color='#e18d4d', s=10)  # 设置颜色为#ff9d6b
                if 'livingCost' in str(data):
                    ax.scatter(x, y2, color='#6fa7f8', s=10)  # 设置颜色为#52c0fb
                # 设置x轴的刻度为整数
                ax.xaxis.set_major_locator(ticker.MaxNLocator(integer=True))
                plt.title('人生财务走势')

                if 'livingCost' in str(data) and 'pensionPrepare' in str(data):
                    ax.fill_between(x, y1, y2, where=(y1 > y2), facecolor='#e18d4d', alpha=0.5)
                    ax.fill_between(x, y1, y2, where=(y1 <= y2), facecolor='#6fa7f8', alpha=0.5)

                # 取消纵坐标科学计数法
                aax = plt.gca()
                aax.ticklabel_format(style='plain')

                # 添加图例
                ax.legend()

                plt.savefig(path)
                return True
        except Exception as ee:
            print(traceback.format_exc())
            return False

def plot_casual_graph(cots):
    G = nx.DiGraph()
    edge_labels = {}
    for cot in cots:
        nodes = cot.split("->")
        for i in range(len(nodes) - 1):
            node_name_i = nodes[i].strip()
            node_name_j = nodes[i + 1].strip()

            G.add_node(node_name_i)
            G.add_node(node_name_j)
            G.add_edge(node_name_i, node_name_j)

    # 绘制因果图
    plt.figure(figsize=(24, 16))
    pos = nx.spring_layout(G)
    nx.draw(G, with_labels=True, node_color='lightblue', node_size=2500, font_size=15)
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=12, rotate=False)
    plt.title('Causal Graph: Age Stage and Insurance Suggestions')
    plt.show()

if __name__ == '__main__':
    data = {
        'summaryVo': {'gap_fund_accumulate': 11488474, 'yearly_income_surplus': 233772, 'res_max_age_fund': 6470000},
        'detailsVos': [{'age': 50, 'livingCost': 325100, 'pensionPrepare': 121992},
                       {'age': 51, 'livingCost': 334853, 'pensionPrepare': 126628},
                       {'age': 52, 'livingCost': 344899, 'pensionPrepare': 131440},
                       {'age': 53, 'livingCost': 355246, 'pensionPrepare': 136434},
                       {'age': 54, 'livingCost': 365903, 'pensionPrepare': 141619},
                       {'age': 55, 'livingCost': 376880, 'pensionPrepare': 147000},
                       {'age': 56, 'livingCost': 388186, 'pensionPrepare': 152586},
                       {'age': 57, 'livingCost': 399832, 'pensionPrepare': 158385},
                       {'age': 58, 'livingCost': 411827, 'pensionPrepare': 164403},
                       {'age': 59, 'livingCost': 424182, 'pensionPrepare': 170651},
                       {'age': 60, 'livingCost': 436907, 'pensionPrepare': 177135},
                       {'age': 61, 'livingCost': 450014, 'pensionPrepare': 183866},
                       {'age': 62, 'livingCost': 463515, 'pensionPrepare': 190853},
                       {'age': 63, 'livingCost': 477420, 'pensionPrepare': 198106},
                       {'age': 64, 'livingCost': 491743, 'pensionPrepare': 205634},
                       {'age': 65, 'livingCost': 506495, 'pensionPrepare': 213448},
                       {'age': 66, 'livingCost': 521690, 'pensionPrepare': 221559},
                       {'age': 67, 'livingCost': 537341, 'pensionPrepare': 229978},
                       {'age': 68, 'livingCost': 553461, 'pensionPrepare': 238717},
                       {'age': 69, 'livingCost': 570065, 'pensionPrepare': 247789},
                       {'age': 70, 'livingCost': 587167, 'pensionPrepare': 257204},
                       {'age': 71, 'livingCost': 604782, 'pensionPrepare': 266978},
                       {'age': 72, 'livingCost': 622925, 'pensionPrepare': 277123},
                       {'age': 73, 'livingCost': 641613, 'pensionPrepare': 287654},
                       {'age': 74, 'livingCost': 660861, 'pensionPrepare': 298585},
                       {'age': 75, 'livingCost': 680687, 'pensionPrepare': 309931},
                       {'age': 76, 'livingCost': 701108, 'pensionPrepare': 321709},
                       {'age': 77, 'livingCost': 722141, 'pensionPrepare': 333934},
                       {'age': 78, 'livingCost': 743805, 'pensionPrepare': 346623},
                       {'age': 79, 'livingCost': 766119, 'pensionPrepare': 359795},
                       {'age': 80, 'livingCost': 789103, 'pensionPrepare': 373467},
                       {'age': 81, 'livingCost': 812776, 'pensionPrepare': 387659},
                       {'age': 82, 'livingCost': 837159, 'pensionPrepare': 402390},
                       {'age': 83, 'livingCost': 862274, 'pensionPrepare': 417680},
                       {'age': 84, 'livingCost': 888142, 'pensionPrepare': 433552},
                       {'age': 85, 'livingCost': 914787, 'pensionPrepare': 450027}],
        'user_info': {'current_age': 30, 'expect_retire_age': 50, 'retire_cost': 15000, 'sex': '女',
                      'have_social_security': '是', 'start_security_num': 8}}
    # data = {'summaryVo': {'legal_retirement_age': 60, 'retire_reserve_fund': 0},
    #         'detailsVos': [{'age': 60, 'pensionPrepare': 125357}, {'age': 61, 'pensionPrepare': 130121},
    #                        {'age': 62, 'pensionPrepare': 135065}, {'age': 63, 'pensionPrepare': 140198},
    #                        {'age': 64, 'pensionPrepare': 145525}, {'age': 65, 'pensionPrepare': 151055},
    #                        {'age': 66, 'pensionPrepare': 156795}, {'age': 67, 'pensionPrepare': 162754},
    #                        {'age': 68, 'pensionPrepare': 168938}, {'age': 69, 'pensionPrepare': 175358},
    #                        {'age': 70, 'pensionPrepare': 182021}, {'age': 71, 'pensionPrepare': 188938},
    #                        {'age': 72, 'pensionPrepare': 196118}, {'age': 73, 'pensionPrepare': 203570},
    #                        {'age': 74, 'pensionPrepare': 211306}, {'age': 75, 'pensionPrepare': 219336},
    #                        {'age': 76, 'pensionPrepare': 227670}, {'age': 77, 'pensionPrepare': 236322},
    #                        {'age': 78, 'pensionPrepare': 245302}, {'age': 79, 'pensionPrepare': 254624},
    #                        {'age': 80, 'pensionPrepare': 264299}, {'age': 81, 'pensionPrepare': 274343},
    #                        {'age': 82, 'pensionPrepare': 284768}, {'age': 83, 'pensionPrepare': 295589},
    #                        {'age': 84, 'pensionPrepare': 306821}, {'age': 85, 'pensionPrepare': 318480}],
    #         'user_info': {'current_age': 40, 'expect_retire_age': 55, 'sex': '男', 'have_social_security': '是',
    #                       'start_security_num': 10}}

    DrawPicture.save_financial_trend(data,
                                     'f0y2KRWddfsffcE0GEew-099-1-output.png')
